﻿using Common;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using Util;

namespace Linguist.Acoustic
{
    /// <summary>
    /// The HMMPool provides the ability to manage units via small integer IDs.  Context Independent units and context
    /// dependent units can be converted to an ID. IDs can be used to quickly retrieve a unit or an hmm associated with the
    /// unit.  This class operates under the constraint that context sizes are exactly one, which is generally only valid for
    /// large vocabulary tasks.
    /// </summary>
    public class HMMPool
    {
        private IAcousticModel model;
        private IUnit[] unitTable;
        private Dictionary<HMMPosition, IHMM[]> hmmTable;
        private int numCIUnits;
        private UnitManager unitManager;


        protected HMMPool(){
        }

        /**
        /// Constructs a HMMPool object.
         *
        /// @param model the model to use for the pool
         */
        public HMMPool(IAcousticModel model, UnitManager unitManager)
        {
            int maxCIUnits = 0;
            this.model = model;
            this.unitManager = unitManager;
            TimerPool.getTimer(this,"Build HMM Pool").start();

            if (model.getLeftContextSize() != 1)
                throw new Exception("LexTreeLinguist: Unsupported left context size");

            if (model.getRightContextSize() != 1)
                throw new Exception("LexTreeLinguist: Unsupported right context size");

            // count CI units:
            IEnumerator<IUnit> i = model.getContextIndependentUnitIterator();
            while (i.MoveNext()) 
            {
                IUnit unit = i.Current;
                //Trace.WriteLine("CI unit " + unit);
                if (unit.getBaseID() > maxCIUnits) 
                {
                    maxCIUnits = unit.getBaseID();
                }
            }

            numCIUnits = maxCIUnits + 1;

            unitTable = new Unit[numCIUnits* numCIUnits* numCIUnits];
            IEnumerator<IHMM> iHMM = model.getHMMIterator();
            while (iHMM.MoveNext()) 
            {
                IHMM hmm = iHMM.Current;
                IUnit unit = hmm.getUnit();
                int id = getID(unit);
                unitTable[id] = unit;

                //Trace.WriteLine("Unit " + unit + " id " + id);

            }

            // build up the hmm table to allow quick access to the hmms
            hmmTable = new Dictionary<HMMPosition, IHMM[]>();
            foreach (HMMPosition position in Enum.GetValues(typeof(HMMPosition))) 
            {
                IHMM[] hmms = new IHMM[unitTable.Length];
                hmmTable.Add(position, hmms);
                for (int j = 1; j < unitTable.Length; j++) 
                {
                    IUnit unit = unitTable[j];
                    if (unit == null) {
                        unit = synthesizeUnit(j);
                    }
                    if (unit != null) {
                        hmms[j] = model.lookupNearestHMM(unit, position, false);
                        Trace.Assert(hmms[j] != null);
                    }
                }
            }
            TimerPool.getTimer(this,"Build HMM Pool").stop();
        }


        public IAcousticModel getModel() 
        {
            return model;
        }


        /**
        /// Given a unit ID, generate a full context dependent unit that will allow us to look for a suitable hmm
         *
        /// @param id the unit id
        /// @return a context dependent unit for the ID
         */
        private IUnit synthesizeUnit(int id) 
        {
            int centralID = getCentralUnitID(id);
            int leftID = getLeftUnitID(id);
            int rightID = getRightUnitID(id);

            if (centralID == 0 || leftID == 0 || rightID == 0) {
                return null;
            }

            IUnit centralUnit = unitTable[centralID];
            IUnit leftUnit = unitTable[leftID];
            IUnit rightUnit = unitTable[rightID];

            Trace.Assert(centralUnit != null);
            Trace.Assert(leftUnit != null);
            Trace.Assert(rightUnit != null);

            IUnit[] lc = new IUnit[1];
            IUnit[] rc = new IUnit[1];
            lc[0] = leftUnit;
            rc[0] = rightUnit;
            LeftRightContext context = LeftRightContext.get(lc, rc);

            IUnit unit = unitManager.getUnit(
                    centralUnit.getName(), centralUnit.isFiller(),
                    context);


            //Trace.WriteLine("Missing " + getUnitNameFromID(id) + " returning " + unit);

            return unit;
        }


        /**
        /// Returns the number of CI units
         *
        /// @return the number of CI Units
         */
        public int getNumCIUnits() 
        {
            return numCIUnits;
        }


        /**
        /// Gets the unit for the given id
         *
        /// @param unitID the id for the unit
        /// @return the unit associated with the ID
         */
        public IUnit getUnit(int unitID) 
        {
            return unitTable[unitID];
        }


        /**
        /// Given a unit id and a position, return the HMM associated with the
        /// unit/position.
         *
        /// @param unitID   the id of the unit
        /// @param position the position within the word
        /// @return the hmm associated with the unit/position
         */
        public IHMM getHMM(int unitID, HMMPosition position) 
        {
            return hmmTable[position][unitID];
        }


        /**
        /// given a unit return its ID
         *
        /// @param unit the unit
        /// @return an ID
         */
        public int getID(IUnit unit) 
        {
            if (unit.isContextDependent()) 
            {
                LeftRightContext context = (LeftRightContext) unit.getContext();
                Trace.Assert(context.getLeftContext().Length == 1);
                Trace.Assert(context.getRightContext().Length == 1);
                return buildID(getSimpleUnitID(unit),
                               getSimpleUnitID(context.getLeftContext()[0]),
                               getSimpleUnitID(context.getRightContext()[0]));
            } 
            else 
            {
                return getSimpleUnitID(unit);
            }
        }


        /**
        /// Returns a context independent ID
         *
        /// @param unit the unit of interest
        /// @return the ID of the central unit (ignoring any context)
         */
        private int getSimpleUnitID(IUnit unit) 
        {
            return unit.getBaseID();
        }


        public Boolean isValidID(int unitID) 
        {
            return unitID >= 0 &&
                   unitID < unitTable.Length &&
                   unitTable[unitID] != null;
        }


        /**
        /// Builds an id from the given unit and its left and right unit ids
         *
        /// @param unitID  the id of the central unit
        /// @param leftID  the id of the left context unit
        /// @param rightID the id of the right context unit
        /// @return the id for the context dependent unit
         */
        public int buildID(int unitID, int leftID, int rightID) 
        {
            // special case ... if the unitID is associated with
            // filler than we have no context ... so use the CI
            // form

            if (unitTable[unitID] == null)
                return -1;

            int id;
            if (unitTable[unitID].isFiller()) {
                id = unitID;
            } else {
                id = unitID* (numCIUnits* numCIUnits)
                        + (leftID* numCIUnits)
                        + rightID;
            }

            Trace.Assert(id < unitTable.Length);

            return id;
        }


        /**
        /// Given a unit id extract the left context unit id
         *
        /// @param id the unit id
        /// @return the unit id of the left context (0 means no left context)
         */
        private int getLeftUnitID(int id) 
        {
            return (id / numCIUnits) % numCIUnits;
        }


        /**
        /// Given a unit id extract the right context unit id
         *
        /// @param id the unit id
        /// @return the unit id of the right context (0 means no right context)
         */
        private int getRightUnitID(int id) 
        {
            return id % numCIUnits;
        }


        /**
        /// Given a unit id extract the central unit id
         *
        /// @param id the unit id
        /// @return the central unit id
         */
        private int getCentralUnitID(int id) 
        {
            return id / (numCIUnits* numCIUnits);
        }


        /**
        /// Given an ID, build up a name for display
         *
        /// @return the name baed on the ID
         */
        private String getUnitNameFromID(int id) 
        {
            int centralID = getCentralUnitID(id);
            int leftID = getLeftUnitID(id);
            int rightID = getRightUnitID(id);

            String cs = unitTable[centralID] == null ? "(" + centralID + ')' :
                    unitTable[centralID].ToString();
            String ls = unitTable[leftID] == null ? ("(" + leftID + ')') :
                    unitTable[leftID].ToString();
            String rs = unitTable[rightID] == null ? "(" + rightID + ')' :
                    unitTable[rightID].ToString();

            return cs + '[' + ls + ',' + rs + ']';
        }

        /**
        /// Retrieves an HMM for a unit in context. If there is no direct match, the
        /// nearest match will be used. Note that we are currently only dealing with,
        /// at most, single unit left and right contexts.
        /// 
        /// @param base
        ///            the base CI unit
        /// @param lc
        ///            the left context
        /// @param rc
        ///            the right context
        /// @param pos
        ///            the position of the base unit within the word
        /// @return the HMM. (This should never return null)
         */
        public IHMM getHMM(IUnit _base, IUnit lc, IUnit rc, HMMPosition pos) 
        {
            int id = -1;
            int bid = getID(_base);
            int lid = getID(lc);
            int rid = getID(rc);

            if (!isValidID(bid)) {
                //Trace.TraceError("Bad HMM Unit: " + _base.getName());
                return null;
            }
            if (!isValidID(lid)) {
                //Trace.TraceError("Bad HMM Unit: " + lc.getName());
                return null;
            }
            if (!isValidID(rid)) {
                //Trace.TraceError("Bad HMM Unit: " + rc.getName());
                return null;
            }
            id = buildID(bid, lid, rid);
            if (id < 0) {
                //Trace.TraceError("Unable to build HMM Unit ID for " + _base.GetType().Name
                //        + " lc=" + lc.getName() + " rc=" + rc.getName());
                return null;
            }
            IHMM hmm = getHMM(id, pos);
            if (hmm == null) {
                //Trace.TraceError("Missing HMM Unit for " + _base.GetType().Name + " lc="
                //        + lc.getName() + " rc=" + rc.getName());
            }

            return hmm;
        }

    

        /** Dumps out info about this pool */
        public void dumpInfo() {
            Trace.TraceInformation("Max CI Units " + numCIUnits);
            Trace.TraceInformation("Unit table size " + unitTable.Length);


            for (int i = 0; i < unitTable.Length; i++) {
                Trace.TraceInformation(i.ToString() + ' ' + unitTable[i]);
            }

        }


        /**
        /// A quick and dirty benchmark to get an idea how long the HMM lookups will take.  This experiment shows that on a
        /// 1GHZ sparc system, the lookup takes a little less than 1uSec.  This is probably fast enough.
         */

        static HMMPosition[] pos = {
                HMMPosition.BEGIN, HMMPosition.END, HMMPosition.SINGLE,
                HMMPosition.INTERNAL};    
    
        static int[] ids = {9206, 9320, 9620, 9865, 14831, 15836};

        void benchmark() {
            int nullCount = 0;
            Trace.WriteLine("benchmarking ...");
            TimerPool.getTimer(this,"hmmPoolBenchmark").start();

            for (int i = 0; i < 1000000; i++) {
                int id = ids[i % ids.Length];
                HMMPosition position = pos[i % pos.Length];
                IHMM hmm = getHMM(id, position);
                if (hmm == null) {
                    nullCount++;
                }
            }
            TimerPool.getTimer(this,"hmmPoolBenchmark").stop();
            Trace.WriteLine("null count " + nullCount);
        }

    }
}
